#include "fermiqcd.h"
#ifdef PARALLEL
ERROR this is not supposed to run parallel
#endif

// required to read the data in
class open_indices {
public:
  mdp_complex c[4][4][3][3];
};

// this stores a 4-quark operator and computes the 2 possible contractions
class op4quark {
public:
  mdp_matrix spin1;
  mdp_matrix color1;
  mdp_matrix spin2;
  mdp_matrix color2;  
  mdp_array<mdp_complex,2> traceAdata;
  mdp_array<mdp_complex,2> traceBdata;

  op4quark(mdp_field<open_indices>& open_prop,
	   mdp_matrix spin1,
	   mdp_matrix color1,
	   mdp_matrix spin2,
	   mdp_matrix color2) {
    this->spin1=spin1;
    this->color1=color1;
    this->spin2=spin2;
    this->color2=color2;
    compute_contractions(open_prop);
  }
  mdp_complex symmetrized_traces(int t1, int t2) {
    // symmetrize results
    return 0.5*(traceAdata(t1,t2)+traceBdata(t1,t2)+
		traceAdata(t2,t1)+traceBdata(t2,t1));
  }

private:
  // Trace{..} Trace{..}
  mdp_complex traceA(mdp_field<open_indices>& open_prop, int t1, int t2) {
    mdp_site x1d(open_prop.lattice());
    mdp_site x2d(open_prop.lattice());
    int nt=open_prop.lattice().size(0);
    x1d.set((t1+nt) % nt);
    x2d.set((t2+nt) % nt);
    mdp_complex c1=0.0,c2=0.0;
    int a,b,i,j;
    for(a=0; a<4; a++)
      for(b=0; b<4; b++)
	for(i=0; i<3; i++) 
	  for(j=0; j<3; j++) {
	    c1+=spin1(b,a)*color1(j,i)*open_prop(x1d).c[a][b][i][j];
	    c2+=spin2(b,a)*color2(j,i)*open_prop(x2d).c[a][b][i][j];
	  }
    return c1*c2;
  }
  // Trace {.. ..}
  mdp_complex traceB(mdp_field<open_indices>& open_prop, int t1, int t2) {
    mdp_site x1d(open_prop.lattice());
    mdp_site x2d(open_prop.lattice());
    int nt=open_prop.lattice().size(0);
    x1d.set((t1+nt) % nt);
    x2d.set((t2+nt) % nt);
    mdp_complex c1=0.0;
    mdp_complex c2=0.0;
    mdp_complex tmp1[4][4][3][3];
    mdp_complex tmp2[4][4][3][3];
    mdp_complex tmp3[4][4][3][3];
    mdp_complex tmp4[4][4][3][3];
    int a,b,c,d,i,j,m,n;
    //
    //  I thinks the following code does...
    //  spin1(b,a)*color1(j,m)*open_prop(x1d).c[c][b][i][j]*
    //  spin2(d,c)*color2(n,i)*open_prop(x2d).c[a][d][m][n];
    //  It could be simplied using Fierz identities but this appear to be fast
    //
    for(a=0; a<4; a++)
      for(c=0; c<4; c++)
	for(i=0; i<3; i++) 
	  for(j=0; j<3; j++) {
	    tmp1[a][c][i][j]=0.0;
	    tmp2[a][c][i][j]=0.0;
	    for(b=0; b<4; b++) {
	      tmp1[a][c][i][j]+=spin1(a,b)*open_prop(x1d).c[b][c][i][j];
	      tmp2[a][c][i][j]+=spin2(a,b)*open_prop(x2d).c[b][c][i][j];
	    }
	  }
    for(a=0; a<4; a++)
      for(c=0; c<4; c++)
	for(i=0; i<3; i++) 
	  for(m=0; m<3; m++) {
	    tmp3[a][c][i][m]=0.0;
	    tmp4[a][c][i][m]=0.0;
	    for(j=0; j<3; j++) {
	      tmp3[a][c][i][m]+=color1(i,j)*tmp1[a][c][j][m];
	      tmp4[a][c][i][m]+=color2(i,j)*tmp2[a][c][j][m];
	    }
	  }
    for(a=0; a<4; a++)
      for(c=0; c<4; c++)
	for(i=0; i<3; i++) 
	  for(m=0; m<3; m++) {
	    c1+=tmp3[a][a][i][i]*tmp4[c][c][m][m];
#ifdef TEST
	    c2+=tmp3[a][c][m][i]*tmp4[a][c][m][i];
#else
	    c2+=tmp3[a][c][m][i]*tmp4[c][a][i][m];
#endif
	  }
    // c1 used as a check. If I output c1 traceA same as traceB. c.v.d.
    return c2;
  }
  void compute_contractions(mdp_field<open_indices>& open_prop) {
    int n=open_prop.lattice().size(0)/2;
    traceAdata.dimension(n,n);
    traceBdata.dimension(n,n);
    for(int t1=0; t1<n; t1++) 
      for(int t2=0; t2<n; t2++) {
	traceAdata(t1,t2)=traceA(open_prop,-t1,t2);
	traceBdata(t1,t2)=traceB(open_prop,-t1,t2);
      }
  }
};

class combo_operator {
public:
  int n;
  mdp_array<mdp_complex,2> output;
  mdp_field<open_indices>* popen_prop;
  combo_operator() {
    popen_prop=0;
  }
  void connect(mdp_field<open_indices>& open_prop) {
    popen_prop=&open_prop;
    n=open_prop.lattice().size(0)/2;
    output.dimension(n,n);
    for(int t1=0; t1<n; t1++) 
      for(int t2=0; t2<n; t2++) {
	output(t1,t2)=0;
      }
  }
  void add_1x1(mdp_complex c, mdp_matrix spin1, mdp_matrix spin2) {
    if(popen_prop==0) { mdp << "must connect before add\n"; exit(1); }
    mdp << "adding 1x1 operator...tracing!\n";    
    op4quark op(*popen_prop,spin1,mdp_identity(3),spin2,mdp_identity(3));
    for(int t1=0; t1<n; t1++) 
      for(int t2=0; t2<n; t2++) {
	output(t1,t2)+=c*op.symmetrized_traces(t1,t2);
      }
  }
  void add_TxT(mdp_complex c, mdp_matrix spin1, mdp_matrix spin2) {
    if(popen_prop==0) { mdp << "must connect before add\n"; exit(1); }
    mdp << "adding TxT operator...tracing!\n";
    for(int a=1; a<9; a++) {
      op4quark op(*popen_prop,spin1,Lambda[a]/2,spin2,Lambda[a]/2);
      for(int t1=0; t1<n; t1++) 
	for(int t2=0; t2<n; t2++) {
	  output(t1,t2)+=c*op.symmetrized_traces(t1,t2);
	}
    }
  }
};

mdp_array<combo_operator,1> 
trace_my_operators(mdp_field<open_indices>& open_prop) {
  mdp_matrix PL=(1-Gamma5)/2;
  mdp_matrix PR=(1+Gamma5)/2;
  mdp_array<combo_operator,1> op(24);
  // OPERATOR 0
  op(0).connect(open_prop);
  op(0).add_1x1(1,PL,PL);
  // OPERATOR 1
  op(1).connect(open_prop);
  op(1).add_1x1(1,PL,PR);
  // OPERATOR 2
  op(2).connect(open_prop);
  op(2).add_1x1(1,PR,PL);
  // OPERATOR 3
  op(3).connect(open_prop);
  op(3).add_1x1(1,PR,PR);
  // OPERATOR 4
  op(4).connect(open_prop);
  for(int mu=0; mu<4; mu++)
    op(4).add_1x1(1,Gamma[mu]*PL,Gamma[mu]*PL);
  // OPERATOR 5
  op(5).connect(open_prop);
  for(int mu=0; mu<4; mu++)
    op(5).add_1x1(1,Gamma[mu]*PL,Gamma[mu]*PR);
  // OPERATOR 6
  op(6).connect(open_prop);
  for(int mu=0; mu<4; mu++)
    op(6).add_1x1(1,Gamma[mu]*PR,Gamma[mu]*PL);
  // OPERATOR 7
  op(7).connect(open_prop);
  for(int mu=0; mu<4; mu++)
    op(7).add_1x1(1,Gamma[mu]*PR,Gamma[mu]*PR);
  // OPERATOR 8
  op(8).connect(open_prop);
  for(int mu=0; mu<4; mu++)
    for(int nu=0; nu<4; nu++)
      if(nu>mu)
	op(8).add_1x1(1,Sigma[mu][nu]*PL,Sigma[mu][nu]*PL);
  // OPERATOR 9
  op(9).connect(open_prop);
  for(int mu=0; mu<4; mu++)
    for(int nu=0; nu<4; nu++)
      if(nu>mu)
	op(9).add_1x1(1,Sigma[mu][nu]*PL,Sigma[mu][nu]*PR);
  // OPERATOR 10
  op(10).connect(open_prop);
  for(int mu=0; mu<4; mu++)
    for(int nu=0; nu<4; nu++)
      if(nu>mu)
	op(10).add_1x1(1,Sigma[mu][nu]*PR,Sigma[mu][nu]*PL);
  // OPERATOR 11
  op(11).connect(open_prop);
  for(int mu=0; mu<4; mu++)
    for(int nu=0; nu<4; nu++)
      if(nu>mu)
	op(11).add_1x1(1,Sigma[mu][nu]*PR,Sigma[mu][nu]*PR);
  // OPERATOR 12
  op(12).connect(open_prop);
  op(12).add_TxT(1,PL,PL);
  // OPERATOR 13
  op(13).connect(open_prop);
  op(13).add_TxT(1,PL,PR);
  // OPERATOR 14
  op(14).connect(open_prop);
  op(14).add_TxT(1,PR,PL);
  // OPERATOR 15
  op(15).connect(open_prop);
  op(15).add_TxT(1,PR,PR);
  // OPERATOR 16
  op(16).connect(open_prop);
  for(int mu=0; mu<4; mu++)
    op(16).add_TxT(1,Gamma[mu]*PL,Gamma[mu]*PL);
  // OPERATOR 17
  op(17).connect(open_prop);
  for(int mu=0; mu<4; mu++)
    op(17).add_TxT(1,Gamma[mu]*PL,Gamma[mu]*PR);
  // OPERATOR 18
  op(18).connect(open_prop);
  for(int mu=0; mu<4; mu++)
    op(18).add_TxT(1,Gamma[mu]*PR,Gamma[mu]*PL);
  // OPERATOR 19
  op(19).connect(open_prop);
  for(int mu=0; mu<4; mu++)
    op(19).add_TxT(1,Gamma[mu]*PR,Gamma[mu]*PR);
  // OPERATOR 20
  op(20).connect(open_prop);
  for(int mu=0; mu<4; mu++)
    for(int nu=0; nu<4; nu++)
      if(nu>mu)
	op(20).add_TxT(1,Sigma[mu][nu]*PL,Sigma[mu][nu]*PL);
  // OPERATOR 21
  op(21).connect(open_prop);
  for(int mu=0; mu<4; mu++)
    for(int nu=0; nu<4; nu++)
      if(nu>mu)
	op(21).add_TxT(1,Sigma[mu][nu]*PL,Sigma[mu][nu]*PL);
  // OPERATOR 22
  op(22).connect(open_prop);
  for(int mu=0; mu<4; mu++)
    for(int nu=0; nu<4; nu++)
      if(nu>mu)
	op(22).add_TxT(1,Sigma[mu][nu]*PL,Sigma[mu][nu]*PL);
  // OPERATOR 23
  op(23).connect(open_prop);
  for(int mu=0; mu<4; mu++)
    for(int nu=0; nu<4; nu++)
      if(nu>mu)
	op(23).add_TxT(1,Sigma[mu][nu]*PL,Sigma[mu][nu]*PL);
  return op;
}

int main(int argc, char** argv) {
  mdp.open_wormholes(argc, argv);
  define_base_matrices("FERMILAB");

  if(mdp.nproc()!=1) {
    mdp << "This cannot run in parallel\n";
    exit(1);
  }

  if(!is_file(argv[1])) {
    mdp << "Unbale to open open_prop file\n";
    exit(1);
  }

  // read the lattice size from the gauge configuration itself
  mdp_field_file_header header=get_info(argv[1]);
  if(header.ndim!=1) {
    mdp << "Invalid input file\n";
    exit(1);
  }

  int* L=header.box;
  int n=3; // number of colors;
  mdp_lattice lattice1d(1,L);
  mdp_site x1d(lattice1d); // site on the projected lattice
  int a,b,c,i,j,k,t,t1,t2,mu,nu; // temp variables
  bool all=false;
  if(argc>2 && strcmp(argv[2],"ALL")==0) all=true;
  mdp_field<open_indices> open_prop(lattice1d); // field for output open prop
  vector<float> c_pi(L[0]); // pion propagator

  // load open propagator
  open_prop.load(argv[1]);

  // make a meson propagator
  for(t=0; t<L[0]; t++) c_pi[t]=0;  
  forallsites(x1d) 
    for(a=0; a<4; a++)
      for(i=0; i<3; i++) {
	b=Gamma5_idx[a];
	c_pi[x1d(0)]+=real(Gamma5_val[a]*open_prop(x1d).c[b][a][i][i]);
      }

  // build all possible 4-quark operators
  mdp_array<combo_operator,1> traces=trace_my_operators(open_prop);

  // FORMATTED OUTPUT HERE //////////////
  char outfile[1024];
  sprintf(outfile, "%s.contractions", argv[1]);
  FILE *fp=fopen(outfile,"w");

  mdp << "OUTPUT t,meson_prop\n";
  for(t=0; t<L[0]; t++) {
    printf("%i, %e\n", t, c_pi[t]);
    fprintf(fp,"%i, %e\n", t, c_pi[t]);
  }
  mdp << "OUTPUT op,t1,t2,real(trace),imag(trace)\n";
  for(int i=0; i<traces.size(); i++) 
    for(int t1=0; t1<traces(i).output.size(0); t1++) 
      for(int t2=0; t2<traces(i).output.size(1); t2++) 
	if(all || t1==t2) {
	  printf("%i, %i, %i, %e, %e\n", i, t1, t2,
		 real(traces(i).output(t1,t2)), 
		 imag(traces(i).output(t1,t2)));
	  fprintf(fp,"%i, %i, %i, %e, %e,\n", i, t1, t2,
		  real(traces(i).output(t1,t2)), 
		  imag(traces(i).output(t1,t2)));
	}  
  fclose(fp);
  
  mdp.close_wormholes();
  return 0;
}

